// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <melon/benchmark.h>
#include <melon/init/init.h>
#include "melon/Uri.h"
#include <pollux/functions/macros.h>
#include <pollux/functions/registerer.h>
#include <pollux/functions/lib/benchmarks/FunctionBenchmarkBase.h>
#include <pollux/functions/prestosql/registration/registration_functions.h>

using namespace kumo::pollux;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::functions;

namespace {

template <typename Fn, typename TOutString, typename TInputString>
MELON_ALWAYS_INLINE bool
runURIFn(Fn func, TOutString& result, TInputString& url) {
  try {
    auto parsedUrl = melon::Uri(url);
    auto datum = (parsedUrl.*func)();
    result.resize(datum.size());
    if (datum.size() != 0) {
      std::strncpy(result.data(), datum.c_str(), datum.size());
    }
  } catch (const std::invalid_argument&) { // thrown if URI is invalid.
    result.resize(0);
  }

  return true;
}

template <typename T>
struct FollyUrlExtractFragmentFunction {
  POLLUX_DEFINE_FUNCTION_TYPES(T);

  // ASCII input always produces ASCII result.
  static constexpr bool is_default_ascii_behavior = true;

  MELON_ALWAYS_INLINE bool call(
      out_type<Varchar>& result,
      const arg_type<Varchar>& urlString) {
    return runURIFn(&melon::Uri::fragment, result, urlString);
  }
};

template <typename T>
struct FollyUrlExtractHostFunction {
  POLLUX_DEFINE_FUNCTION_TYPES(T);

  // ASCII input always produces ASCII result.
  static constexpr bool is_default_ascii_behavior = true;

  MELON_ALWAYS_INLINE bool call(
      out_type<Varchar>& result,
      const arg_type<Varchar>& urlString) {
    return runURIFn(&melon::Uri::host, result, urlString);
  }
};

template <typename T>
struct FollyUrlExtractPathFunction {
  POLLUX_DEFINE_FUNCTION_TYPES(T);

  // ASCII input always produces ASCII result.
  static constexpr bool is_default_ascii_behavior = true;

  MELON_ALWAYS_INLINE bool call(
      out_type<Varchar>& result,
      const arg_type<Varchar>& urlString) {
    return runURIFn(&melon::Uri::path, result, urlString);
  }
};

template <typename T>
struct FollyUrlExtractQueryFunction {
  POLLUX_DEFINE_FUNCTION_TYPES(T);

  // ASCII input always produces ASCII result.
  static constexpr bool is_default_ascii_behavior = true;

  MELON_ALWAYS_INLINE bool call(
      out_type<Varchar>& result,
      const arg_type<Varchar>& urlString) {
    return runURIFn(&melon::Uri::query, result, urlString);
  }
};

template <typename T>
struct FollyUrlExtractProtocolFunction {
  POLLUX_DEFINE_FUNCTION_TYPES(T);

  // ASCII input always produces ASCII result.
  static constexpr bool is_default_ascii_behavior = true;

  MELON_ALWAYS_INLINE bool call(
      out_type<Varchar>& result,
      const arg_type<Varchar>& urlString) {
    return runURIFn(&melon::Uri::scheme, result, urlString);
  }
};

template <typename T>
struct FollyUrlExtractPortFunction {
  POLLUX_DEFINE_FUNCTION_TYPES(T);
  // ASCII input always produces ASCII result.
  static constexpr bool is_default_ascii_behavior = true;

  MELON_ALWAYS_INLINE bool call(int64_t& result, const arg_type<Varchar>& url) {
    try {
      auto parsedUrl = melon::Uri(url);
      result = parsedUrl.port();
    } catch (const melon::ConversionError&) {
      return false;
    } catch (const std::invalid_argument&) {
      return false;
    }
    return true;
  }
};

template <typename T>
struct FollyUrlExtractParameterFunction {
  POLLUX_DEFINE_FUNCTION_TYPES(T);
  // ASCII input always produces ASCII result.
  static constexpr bool is_default_ascii_behavior = true;

  MELON_ALWAYS_INLINE bool call(
      out_type<Varchar>& result,
      const arg_type<Varchar>& url,
      const arg_type<Varchar>& param) {
    try {
      auto parsedUrl = melon::Uri(url);
      auto params = parsedUrl.getQueryParams();

      for (const auto& pair : params) {
        if (std::string(param) == pair.first) {
          result.resize(pair.second.length());
          strncpy(result.data(), pair.second.c_str(), pair.second.length());
        }
      }
    } catch (const std::invalid_argument&) {
      return false;
    }
    return false;
  }
};

class UrlBenchmark : public functions::test::FunctionBenchmarkBase {
 public:
  UrlBenchmark() : FunctionBenchmarkBase() {
    functions::prestosql::registerURLFunctions();

    // Register melon based implementations.
    register_function<FollyUrlExtractFragmentFunction, Varchar, Varchar>(
        {"melon_url_extract_fragment"});
    register_function<FollyUrlExtractHostFunction, Varchar, Varchar>(
        {"melon_url_extract_host"});
    register_function<FollyUrlExtractPathFunction, Varchar, Varchar>(
        {"melon_url_extract_path"});
    register_function<FollyUrlExtractQueryFunction, Varchar, Varchar>(
        {"melon_url_extract_query"});
    register_function<FollyUrlExtractProtocolFunction, Varchar, Varchar>(
        {"melon_url_extract_protocol"});
    register_function<FollyUrlExtractPortFunction, int64_t, Varchar>(
        {"melon_url_extract_port"});
    register_function<
        FollyUrlExtractParameterFunction,
        Varchar,
        Varchar,
        Varchar>({"melon_url_extract_parameter"});
  }

  void runUrlExtract(const std::string& fnName, bool isParameter = false) {
    melon::BenchmarkSuspender suspender;

    size_t size = 1000;
    std::string url;
    auto vectorUrls = vectorMaker_.flat_vector<StringView>(
        size,
        [&](auto row) {
          // construct some pseudo random url
          url = fmt::format(
              "http://somehost{}.com:8080/somepath{}/p.php?k1={}#Refi",
              row,
              row % 2,
              row % 3);
          return StringView(url);
        },
        nullptr);
    auto constVector =
        BaseVector::create_constant(VARCHAR(), "k1", size, pool());
    auto row_vector = isParameter
        ? vectorMaker_.row_vector({vectorUrls, constVector})
        : vectorMaker_.row_vector({vectorUrls});

    auto queryString = isParameter ? "{}(c0, c1)" : "{}(c0)";
    auto exprSet = compileExpression(
        fmt::format(fmt::runtime(queryString), fnName), row_vector->type());

    suspender.dismiss();

    doRun(exprSet, row_vector);
  }

  void doRun(ExprSet& exprSet, const RowVectorPtr& row_vector) {
    uint32_t cnt = 0;
    for (auto i = 0; i < 100; i++) {
      cnt += evaluate(exprSet, row_vector)->size();
    }
    melon::doNotOptimizeAway(cnt);
  }
};

BENCHMARK(melon_fragment) {
  UrlBenchmark benchmark;
  benchmark.runUrlExtract("melon_url_extract_fragment");
}

BENCHMARK_RELATIVE(pollux_fragment) {
  UrlBenchmark benchmark;
  benchmark.runUrlExtract("url_extract_fragment");
}

BENCHMARK(melon_host) {
  UrlBenchmark benchmark;
  benchmark.runUrlExtract("melon_url_extract_host");
}

BENCHMARK_RELATIVE(pollux_host) {
  UrlBenchmark benchmark;
  benchmark.runUrlExtract("url_extract_host");
}

BENCHMARK(melon_path) {
  UrlBenchmark benchmark;
  benchmark.runUrlExtract("melon_url_extract_path");
}

BENCHMARK_RELATIVE(pollux_path) {
  UrlBenchmark benchmark;
  benchmark.runUrlExtract("url_extract_path");
}

BENCHMARK(melon_query) {
  UrlBenchmark benchmark;
  benchmark.runUrlExtract("melon_url_extract_query");
}

BENCHMARK_RELATIVE(pollux_query) {
  UrlBenchmark benchmark;
  benchmark.runUrlExtract("url_extract_query");
}

BENCHMARK(melon_port) {
  UrlBenchmark benchmark;
  benchmark.runUrlExtract("melon_url_extract_port");
}

BENCHMARK_RELATIVE(pollux_port) {
  UrlBenchmark benchmark;
  benchmark.runUrlExtract("url_extract_port");
}

BENCHMARK(melon_protocol) {
  UrlBenchmark benchmark;
  benchmark.runUrlExtract("melon_url_extract_protocol");
}

BENCHMARK_RELATIVE(pollux_protocol) {
  UrlBenchmark benchmark;
  benchmark.runUrlExtract("url_extract_protocol");
}

BENCHMARK(melon_param) {
  UrlBenchmark benchmark;
  benchmark.runUrlExtract("melon_url_extract_parameter", true);
}

BENCHMARK_RELATIVE(pollux_param) {
  UrlBenchmark benchmark;
  benchmark.runUrlExtract("url_extract_parameter", true);
}

} // namespace

int main(int argc, char** argv) {
  melon::Init init{&argc, &argv};

  melon::runBenchmarks();
  return 0;
}
